import numpy as np
import matplotlib.pyplot as plt

n = 64
h = 1/n
w1 = 1
w2 = 1/3
w3 = 1/2
w4 = 2/3

k = np.linspace(0, 64, 10000)

# 计算 y = sin(k * x * pi)
y1 = 1 - 2*w1*(np.sin(k * np.pi / (2 * n))**2)
y2 = 1 - 2*w2*(np.sin(k * np.pi / (2 * n))**2)
y3 = 1 - 2*w3*(np.sin(k * np.pi / (2 * n))**2)
y4 = 1 - 2*w4*(np.sin(k * np.pi / (2 * n))**2)

# 创建一个图形对象，并在其中创建两个子图
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 6))

# 第一个子图
ax1.plot(k, y1, 'r', label=f'w1 = 1')
ax1.plot(k, y2, 'g', label=f'w2 = 1/3')
ax1.plot(k, y3, 'b', label=f'w3 = 1/2')
ax1.plot(k, y4, 'y', label=f'w4 = 2/3')
ax1.set_xlabel('k')
ax1.set_ylabel('$1 - 2w*sin^2(kπ/2n)$')
ax1.set_title('E9.18: Plot 1')
ax1.grid(True)
ax1.legend()

# 第二个子图
w = np.linspace(0, 1, 10000)
p = [max(np.abs(1 - 2 * w1 * (np.sin(np.pi * t / (2 * n)))**2) for t in range(1, n)) for w1 in w]

ax2.plot(w, p, label=f'p($T_w$)')
ax2.set_xlabel('w')
ax2.set_ylabel('p($T_w$)')
ax2.set_title('E9.18: Plot 2')
ax2.grid(True)
ax2.legend()

# 调整布局并显示图像
plt.tight_layout()
plt.savefig("./E9_18.jpg")